"""
# @Time    : 2021/6/30 10:07 下午
# @Author  : hezhiqiang
# @Email   : tinyzqh@163.com
# @File    : train.py
"""

# !/usr/bin/env python
import sys
import os
import socket
import setproctitle
import numpy as np
from pathlib import Path
import torch

# 添加CUDA库路径
os.environ['PATH'] = os.environ['PATH'] + ';C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.7\\bin'

# Get the parent directory of the current file
parent_dir = os.path.abspath(os.path.join(os.getcwd(), "."))

# Append the parent directory to sys.path, otherwise the following import will fail
sys.path.append(parent_dir)

# 获取项目根目录路径
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(root_dir)

# 修改导入语句
try:
    from config import get_config
except ModuleNotFoundError:
    # 尝试从项目根目录导入
    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
    from config import get_config
    
from envs.env_wrappers import DummyVecEnv

"""MPE环境的训练脚本。"""


def make_train_env(all_args):
    """
    创建训练环境
    
    参数:
        all_args: 所有参数
        
    返回:
        训练环境实例
    """
    def get_env_fn(rank):
        """
        获取环境函数
        
        参数:
            rank: 环境序号
            
        返回:
            初始化环境的函数
        """
        def init_env():
            """
            初始化环境
            
            返回:
                环境实例
            """
            # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.

            # from envs.env_continuous import ContinuousActionEnv

            # env = ContinuousActionEnv()  # 创建连续动作环境

            # from envs.env_discrete import DiscreteActionEnv

            # env = DiscreteActionEnv()  # 创建离散动作环境

            # env.seed(all_args.seed + rank * 1000)  # 设置环境种子
            # return env


            # 使用UAV CPP离散动作环境
            from envs.env_uav_cpp import UAVCPPDiscreteActionEnv
            
            env = UAVCPPDiscreteActionEnv()  # 创建UAV CPP离散动作环境
            
            env.seed(all_args.seed + rank * 1000)  # 设置环境种子
            return env

        return init_env

    return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)])  # 返回向量化环境


def make_eval_env(all_args):
    """
    创建评估环境
    
    参数:
        all_args: 所有参数
        
    返回:
        评估环境实例
    """
    def get_env_fn(rank):
        """
        获取环境函数
        
        参数:
            rank: 环境序号
            
        返回:
            初始化环境的函数
        """
        def init_env():
            """
            初始化环境
            
            返回:
                环境实例
            """
            # TODO 注意注意，这里选择连续还是离散可以选择注释上面两行，或者下面两行。
            # TODO Important, here you can choose continuous or discrete action space by uncommenting the above two lines or the below two lines.
            from envs.env_continuous import ContinuousActionEnv

            env = ContinuousActionEnv()  # 创建连续动作环境
            # from envs.env_discrete import DiscreteActionEnv
            # env = DiscreteActionEnv()  # 创建离散动作环境
            env.seed(all_args.seed + rank * 1000)  # 设置环境种子
            return env

        return init_env

    return DummyVecEnv([get_env_fn(i) for i in range(all_args.n_rollout_threads)])  # 返回向量化环境


def parse_args(args, parser):
    """
    解析命令行参数
    
    参数:
        args: 命令行参数
        parser: 参数解析器
        
    返回:
        解析后的参数
    """
    

    all_args = parser.parse_known_args(args)[0]  # 解析参数

    return all_args


def main(args):
    """
    主函数
    
    参数:
        args: 命令行参数
    """
    parser = get_config()  # 获取配置
    all_args = parse_args(args, parser)  # 解析参数

    # 检查算法和策略设置
    if all_args.algorithm_name == "rmappo":
        assert all_args.use_recurrent_policy or all_args.use_naive_recurrent_policy, "检查循环策略！"
    elif all_args.algorithm_name == "mappo":
        assert (
            all_args.use_recurrent_policy == False and all_args.use_naive_recurrent_policy == False
        ), "检查循环策略！"
    else:
        raise NotImplementedError  # 不支持的算法

    # 检查共享策略和场景设置
    assert (
        all_args.share_policy == True and all_args.scenario_name == "simple_speaker_listener"
    ) == False, "simple_speaker_listener场景不能使用共享策略。请检查config.py。"

    # CUDA设置
    if all_args.cuda and torch.cuda.is_available():
        print("选择使用GPU...")
        device = torch.device("cuda:0")  # 使用GPU
        torch.set_num_threads(all_args.n_training_threads)  # 设置线程数
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False  # 关闭benchmark
            torch.backends.cudnn.deterministic = True  # 设置确定性
    else:
        print("选择使用CPU...")
        device = torch.device("cpu")  # 使用CPU
        torch.set_num_threads(all_args.n_training_threads)  # 设置线程数

    # 运行目录
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )  # 创建运行目录路径
    if not run_dir.exists():
        os.makedirs(str(run_dir))  # 创建运行目录

    # 确定当前运行编号
    if not run_dir.exists():
        curr_run = "run1"  # 第一次运行
    else:
        exst_run_nums = [
            int(str(folder.name).split("run")[1])
            for folder in run_dir.iterdir()
            if str(folder.name).startswith("run")
        ]  # 获取已存在的运行编号
        if len(exst_run_nums) == 0:
            curr_run = "run1"  # 第一次运行
        else:
            curr_run = "run%i" % (max(exst_run_nums) + 1)  # 下一次运行
    run_dir = run_dir / curr_run  # 更新运行目录
    if not run_dir.exists():
        os.makedirs(str(run_dir))  # 创建运行目录

    # 设置进程标题
    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # 设置随机种子
    torch.manual_seed(all_args.seed)  # PyTorch种子
    torch.cuda.manual_seed_all(all_args.seed)  # CUDA种子
    np.random.seed(all_args.seed)  # NumPy种子

    # 初始化环境
    envs = make_train_env(all_args)  # 创建训练环境
    eval_envs = make_eval_env(all_args) if all_args.use_eval else None  # 创建评估环境
    num_agents = all_args.num_agents  # 智能体数量

    # 配置
    config = {
        "all_args": all_args,  # 所有参数
        "envs": envs,  # 训练环境
        "eval_envs": eval_envs,  # 评估环境
        "num_agents": num_agents,  # 智能体数量
        "device": device,  # 设备
        "run_dir": run_dir,  # 运行目录
    }

    # 运行实验
    if all_args.share_policy:
        from runner.shared.env_runner import EnvRunner as Runner  # 共享策略
    else:
        from runner.separated.env_runner import EnvRunner as Runner  # 分离策略

    runner = Runner(config)  # 创建运行器
    runner.run()  # 运行训练

    # 后处理
    envs.close()  # 关闭训练环境
    if all_args.use_eval and eval_envs is not envs:
        eval_envs.close()  # 关闭评估环境

    runner.writter.export_scalars_to_json(str(runner.log_dir + "/summary.json"))  # 导出标量
    runner.writter.close()  # 关闭写入器


if __name__ == "__main__":
    main(sys.argv[1:])  # 运行主函数
